# 基于tools.py 计算范数
import tools
import numpy as np
from basix.ufl import element
from dolfinx import default_real_type, fem


# 适用于向量和标量的 L2范数计算
def get_L2_err(msh, dof_map, num_cells, u_true_func, u_num, func_space, order):
    """ 计算L2误差 
    Args:
        msh: 基于fenicsx定义的网格
        dof_map (_type_): 自由度索引表
        num_cells (_type_): 单元id总数
        u_true_func (_type_): 真解的函数
        u_num (_type_): 变量数值解
        func_space: 基函数空间
        order: 有限元阶数(eg: P1元取order=1)
        
    目前只处理P1 P2元 三角形单元
    u_true_func(x)接口 x[0]存储x坐标 x[1]存储y坐标 
    """
    G = tools.GaussIntegrate(order=9)
    if order == 1:
        P1 = element("Lagrange", msh.basix_cell(), degree=1, dtype=default_real_type)
        basis_values = P1.tabulate(0, G.points)[0]
        num_basis_func = 3
        
    elif order == 2:  # 矢量元 基于标量进行计算
        P2 = element("Lagrange", msh.basix_cell(), degree=2, dtype=default_real_type)
        basis_values = P2.tabulate(0, G.points)[0]
        num_basis_func = 6
    else:
        raise NotImplementedError("只支持P1和P2元")
    
    L2_err = 0.0  # L2误差平方和
    global_dofs = func_space.tabulate_dof_coordinates()  # 网格所有节点坐标 按dof索引存储
    num_gs_points = basis_values.shape[0]  # 高斯点数量

    for cell_idx in range(num_cells):
        dof_idx = dof_map.cell_dofs(cell_idx)  # 获取自由度索引
        local_nodes = global_dofs[dof_idx][:, :-1]  # 自由度节点的全局坐标
        local_vertices = local_nodes[:3]  # 取前3个节点作为顶点
        gs_phy_points = G.to_physical_gauss_points(*local_vertices)
        bdt = tools.BasisDerivativeTransformer(local_vertices)

        if func_space.element.space_dimension / num_basis_func == 1:  # 标量
            u_nodes_num = u_num[dof_idx]  # 数值解的节点值
            u_gs_true = u_true_func(gs_phy_points.T)  # 真解在高斯点的值
            u_gs_num = np.zeros_like(u_gs_true)
            for alpha in range(num_gs_points):
                # 计算高斯点的物理坐标
                u_gs_num[alpha] = np.dot(basis_values[alpha], u_nodes_num)
            r1 = (u_gs_num - u_gs_true) ** 2  # 误差平方
            L2_err += np.dot(r1, G.weights) * bdt.detJ
        elif func_space.element.space_dimension / num_basis_func == 2:  # 二维矢量
            # 矢量分量的计算
            u_x_nodes_num = u_num[2 * dof_idx]
            u_y_nodes_num = u_num[2 * dof_idx + 1]
            u_x_gs_true, u_y_gs_true = u_true_func(gs_phy_points.T)  # 真解
            u_x_gs_num, u_y_gs_num = np.zeros_like(u_x_gs_true), np.zeros_like(u_y_gs_true)
            for alpha in range(num_gs_points):
                u_x_gs_num[alpha] = np.dot(basis_values[alpha], u_x_nodes_num)
                u_y_gs_num[alpha] = np.dot(basis_values[alpha], u_y_nodes_num)
            r1 = (u_x_gs_num - u_x_gs_true) ** 2 + (u_y_gs_num - u_y_gs_true) ** 2
            L2_err += np.dot(r1, G.weights) * bdt.detJ
        else:
            raise NotImplementedError("只支持标量和矢量的L2误差计算")
    return np.sqrt(L2_err)  # 返回L2误差平方根


def get_H1_err(msh, dof_map, num_cells, u_true_func_diff_ls, u_num, func_space, order):
    """ 计算H1误差 
    Args:
        msh: 基于fenicsx定义的网格
        dof_map (_type_): 自由度索引表
        num_cells (_type_): 单元id总数
        u_true_func_diff_ls (list): 真解导数的函数列表
        u_num (_type_): 变量数值解
        func_space: 基函数空间
        order: 有限元阶数(eg: P1元取order=1)
        
    目前只处理P1 P2元 三角形单元
    真解导数格式为 [du/dx, du/dy] 或 [(du/dx, du/dy), (dv/dx, dv/dy)]
    使用(x, y)作为函数参数
    """
    G = tools.GaussIntegrate(order=9)
    
    if order == 1:
        P1 = element("Lagrange", msh.basix_cell(), degree=1, dtype=default_real_type)
        basis_values, basis_values_d_xi, basis_values_d_eta = P1.tabulate(1, G.points)
        num_basis_func = 3
        
    elif order == 2:  # 矢量元 基于标量进行计算
        P2 = element("Lagrange", msh.basix_cell(), degree=2, dtype=default_real_type)
        basis_values, basis_values_d_xi, basis_values_d_eta = P2.tabulate(1, G.points)
        num_basis_func = 6
    else:
        raise NotImplementedError("只支持P1和P2元")
    assert 2 * func_space.element.space_dimension / num_basis_func == len(u_true_func_diff_ls), "真解导数函数列表长度不匹配"
    
    H1_err = 0.0  
    global_dofs = func_space.tabulate_dof_coordinates()  # 网格所有节点坐标 按dof索引存储
    num_gs_points = basis_values.shape[0]  # 高斯点数量

    for cell_idx in range(num_cells):
        dof_idx = dof_map.cell_dofs(cell_idx)  # 获取自由度索引
        local_nodes = global_dofs[dof_idx][:, :-1]  # 自由度节点的全局坐标
        local_vertices = local_nodes[:3]  # 取前3个节点作为顶点
        gs_phy_points = G.to_physical_gauss_points(*local_vertices)
        gs_phy_x = gs_phy_points[:, 0]  # 高斯点的x坐标
        gs_phy_y = gs_phy_points[:, 1]  # 高斯点的y坐标
        bdt = tools.BasisDerivativeTransformer(local_vertices)
        D = 1 if tools.is_counter_clock(local_vertices)  else -1  #  顺时针需要取正数
        local_basis = bdt.local(basis_values)  # 基函数在局部坐标系下的值
        if func_space.element.space_dimension / num_basis_func == 1:  # 标量
            # 读取导数的定义
            u_true_dx, u_true_dy = u_true_func_diff_ls  
            u_nodes_num = u_num[dof_idx]  # 数值解的节点值
            u_gs_true_dx = u_true_dx(gs_phy_x, gs_phy_y)  # 真解在x方向导数高斯点的值
            u_gs_true_dy = u_true_dy(gs_phy_x, gs_phy_y)  # 真解在y方向导数高斯点的值
            
            local_basis_dx = bdt.local_dx1(basis_values_d_xi, basis_values_d_eta)
            local_basis_dy = bdt.local_dy1(basis_values_d_xi, basis_values_d_eta)
            u_num_dx, u_num_dy = np.zeros_like(u_gs_true_dx), np.zeros_like(u_gs_true_dy)  # 数值解在高斯点的导数值
        
            for i in range(num_gs_points):  
                u_num_dx[i] = np.dot(u_nodes_num, local_basis_dx[i]) * D
                u_num_dy[i] = np.dot(u_nodes_num, local_basis_dy[i]) * D
            r1 = (u_num_dx - u_gs_true_dx) ** 2 + (u_num_dy - u_gs_true_dy) ** 2
            H1_err += np.dot(r1, G.weights) * bdt.detJ
            
        elif func_space.element.space_dimension / num_basis_func == 2:  # 二维矢量
            # 矢量分量的计算 u = (u1, u2)
            u1_true_dx, u1_true_dy, u2_true_dx, u2_true_dy = u_true_func_diff_ls 
            u1_nodes_num, u2_nodes_num = u_num[2 * dof_idx], u_num[2 * dof_idx + 1]  # 数值解的节点值
            u1_gs_true_dx = u1_true_dx(gs_phy_x, gs_phy_y)  # u1在x方向导数 高斯点的值
            u1_gs_true_dy = u1_true_dy(gs_phy_x, gs_phy_y)  # u1在y方向导数 高斯点的值
            u2_gs_true_dx = u2_true_dx(gs_phy_x, gs_phy_y)  # u2在x方向导数 高斯点的值
            u2_gs_true_dy = u2_true_dy(gs_phy_x, gs_phy_y)  # u2在y方向导数 高斯点的值
            
            local_basis_dx = bdt.local_dx1(basis_values_d_xi, basis_values_d_eta)
            local_basis_dy = bdt.local_dy1(basis_values_d_xi, basis_values_d_eta)
            
            u1_num_dx, u1_num_dy = np.zeros_like(u1_gs_true_dx), np.zeros_like(u1_gs_true_dy)  # u1数值解在高斯点的导数值
            u2_num_dx, u2_num_dy = np.zeros_like(u2_gs_true_dx), np.zeros_like(u2_gs_true_dy)  # u2数值解在高斯点的导数值
            
            for i in range(num_gs_points):
                u1_num_dx[i] = np.dot(u1_nodes_num, local_basis_dx[i]) * D
                u1_num_dy[i] = np.dot(u1_nodes_num, local_basis_dy[i]) * D
                u2_num_dx[i] = np.dot(u2_nodes_num, local_basis_dx[i]) * D
                u2_num_dy[i] = np.dot(u2_nodes_num, local_basis_dy[i]) * D
            r1 = (u1_num_dx - u1_gs_true_dx) ** 2 + (u1_num_dy - u1_gs_true_dy) ** 2
            r2 = (u2_num_dx - u2_gs_true_dx) ** 2 + (u2_num_dy - u2_gs_true_dy) ** 2
            H1_err += np.dot(r1 + r2, G.weights) * bdt.detJ
        else:
            raise NotImplementedError("只支持标量和矢量的H1误差计算")
    return np.sqrt(H1_err)  # 返回H1误差平方根


def get_max_node_err(u_true_func, u_num, func_space):
    """ 计算最大节点误差
    Args:
        u_true_func (_type_): 真解的函数
        u_num (_type_): 变量数值解
        func_space: 基函数空间
    目前只处理P1 P2元 三角形单元
    u_true_func(x)接口 x[0]存储x坐标 x[1]存储y坐标 
    """
    u_true = fem.Function(func_space)
    u_true.interpolate(u_true_func)  # 插值真解到函数空间
    return np.max(np.abs(u_true.x.array - u_num))  # 返回最大误差


def interpolate(basis_values, u_nodes_num):
    """ 插值计算区域内一个点的数值解 
    Args:
        basis_values: 基函数在需计算点的值
        u_nodes_num: 数值解的节点值
    Returns:
        uh: 计算点的插值结果
    """
    assert len(basis_values) == len(u_nodes_num), "基函数值和节点数值解长度不匹配"
    num_basis_func = len(basis_values) # 基函数数量
    return np.dot(basis_values, u_nodes_num)  # 返回插值结果


def interpolate_vec(basis_values, u_nodes_num):
    """ 插值计算区域内 向量的数值解 
    Args:
        basis_values: 基函数在需计算的值 形状为 (num_gs_points, num_basis_func)
        u_nodes_num: 数值解的节点值
    Returns:
        uh: 计算点的插值结果
    """
    assert basis_values.shape[1] == len(u_nodes_num), "基函数值和节点数值解长度不匹配"
    num_basis_func = len(basis_values) # 基函数数量
    return np.dot(basis_values, u_nodes_num)  # 返回插值结果